/**
 * Copyright (c) 2010
 * Bert Young. UESTC. 
 */
#include <errno.h>
#include <unistd.h>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <netinet/tcp.h>
#include <poll.h>
#include "SocketConnection.h"

#if 0
int SocketConnection::READ_TIMEOUT;
int SocketConnection::WRITE_TIMEOUT;
int SocketConnection::SIZE_MASK;
int SocketConnection::MAX_DATASIZE;
int SocketConnection::HEAD;
int SocketConnection::COMPLETE_READ;
int SocketConnection::COMPLETE_WRITE;
int SocketConnection::TIMEOUT;
int SocketConnection::ERROR;
int SocketConnection::EOFSOCKET;
#endif

SocketConnection::SocketConnection(int localfd, const struct sockaddr_in * remote)
{
	this->bitmask = COMPLETE_READ;//The connectio is new, so kernel recv buffer is empty, ie, unavailable

	this->localSock = localfd;
	
	bzero(&this->localAddr, sizeof(struct sockaddr_in));
	bzero(&this->remoteAddr, sizeof(struct sockaddr_in));
	socklen_t len = sizeof(struct sockaddr);
	if ( remote == NULL )
		::getpeername(this->localSock, (struct sockaddr *)&this->remoteAddr, &len);
	else
		bcopy(remote, &this->remoteAddr, sizeof(struct sockaddr_in));
	::getsockname( this->localSock, (struct sockaddr *)&this->localAddr, &len);

	int flag = ::fcntl(this->localSock, F_GETFD, 0); 
	::fcntl(this->localSock, F_SETFD, flag | O_NONBLOCK | FD_CLOEXEC);

	int nodelay = 1;
	::setsockopt(this->localSock , IPPROTO_TCP, TCP_NODELAY, (void *)&nodelay, sizeof(int) );
}

SocketConnection::~SocketConnection()
{
	TEMP_FAILURE_RETRY( ::shutdown(localSock, SHUT_RDWR) );
	TEMP_FAILURE_RETRY( ::close(this->localSock) );
}

inline int SocketConnection::waitForRead(int waittime)
{
	struct pollfd st = {this->localSock, POLLIN | POLLPRI, 0}; 
	int ret = TEMP_FAILURE_RETRY( ::poll(&st, 1, waittime ) );
	if ( ret == 1 && !(st.revents & POLLIN ) )
		return ERROR;
	else if ( ret == -1 && EAGAIN == errno )
		return TIMEOUT;

	return ret;
}

inline int SocketConnection::waitForWrite(int waittime)
{
	struct pollfd st = {this->localSock, POLLOUT | POLLPRI, 0}; 
	int ret = TEMP_FAILURE_RETRY( ::poll(&st, 1, waittime ) );
	if ( ret == 1 && !(st.revents & POLLOUT ) )
		return ERROR;
	else if ( ret == ERROR && EAGAIN == errno )
		return TIMEOUT;

	return ret;
}

template <typename BUFFER>
inline unsigned int SocketConnection::packPacket(const void * data, BUFFER & dest, unsigned int size)
{
	dest.assureSpace(size + HEAD);
	bcopy(data, dest.writeAddr()+HEAD, size);
	*(unsigned int *)dest.writeAddr() = size;
	dest.adjustWritePtr(size + HEAD);

	return size + HEAD;
}

inline unsigned int SocketConnection::unpackPacket(const unsigned char * in, void * dest, unsigned int wholesize)
{
	bcopy( in+HEAD, dest, wholesize - HEAD );
	return wholesize - HEAD;
}
#if 0
// from recvbuf to our cmd, out should be big enough
int SocketConnection::retrieveOnePacket(void * out )
{
	int nBytes = 0;
	if ( recvBuf.sizeForRead() >= HEAD )
	{
		nBytes = *(unsigned int *)recvBuf.readAddr();
		if ( nBytes + HEAD <= recvBuf.sizeForRead() )
		{
			unpackPacket(recvBuf.readAddr(), out, nBytes+HEAD);
			recvbuf.adjustReadPtr(nBytes + HEAD);
		}
		else
			nBytes = 0;
	}
	return nBytes;
}
#endif
//int recvToBuf();//::recv
//will set mark for kernel buffer
//will try poll once , will get as much data into user buffer
int SocketConnection::recvFromKernel()
{
	int ret = 0;
	if ( bitmask & COMPLETE_READ )
	{
		bitmask &= ~COMPLETE_READ;
		goto pollForRead;
	}

	recvBuf.assureSpace(MAX_DATASIZE);
	ret = TEMP_FAILURE_RETRY( ::recv(this->localSock, recvBuf.writeAddr(), recvBuf.sizeForWrite(), MSG_NOSIGNAL) );
	if ( ERROR == ret && EAGAIN == errno )
	{
pollForRead:
		ret = waitForRead();
		if ( ret > 0 )
			ret = TEMP_FAILURE_RETRY( ::recv(this->localSock, recvBuf.writeAddr(), recvBuf.sizeForWrite(), MSG_NOSIGNAL) );
		else
			return ret;
	}

	if ( ret > 0 )
	{
		if ( static_cast<unsigned int>(ret) < recvBuf.sizeForWrite() )
			bitmask |= COMPLETE_READ;
		recvBuf.adjustWritePtr(ret);
	}
	return (0 == ret) ? EOFSOCKET : ret;
}

int SocketConnection::recv()
{
	recvBuf.assureSpace(MAX_DATASIZE);
	int ret = TEMP_FAILURE_RETRY( ::recv(this->localSock, recvBuf.writeAddr(), recvBuf.sizeForWrite(), MSG_NOSIGNAL) );
	if ( ret == ERROR && errno == EAGAIN )
		return 0;

	if ( ret > 0 )
	{
		if ( static_cast<unsigned int>(ret) < recvBuf.sizeForWrite() )
			bitmask |= COMPLETE_READ;
		recvBuf.adjustWritePtr(ret);
	}
	return (0 == ret) ? EOFSOCKET : ret;

}
//must get a packet; recv a packet until success or die
int SocketConnection::recvPacket(void * dest, bool wait /*true*/ )
{
	int nBytes = 0;
	while ( 0 == (nBytes = recvPacketNoPoll(dest)) )
	{
		nBytes = recvFromKernel();
		if ( nBytes < 0 )
			return ERROR;
		if ( nBytes == 0 && !wait )
			break;
	}
	return nBytes;
}

//recv to cmd no poll, recv a packet or fail; try get a packet
int SocketConnection::recvPacketNoPoll(void * dest)
{
	int nBytes = 0;
	if ( recvBuf.sizeForRead() >= HEAD )
	{
		nBytes = *(unsigned int *)recvBuf.readAddr();
		if ( nBytes + HEAD <= recvBuf.sizeForRead() )
		{
			unpackPacket(recvBuf.readAddr(), dest, nBytes+HEAD);
			recvBuf.adjustReadPtr(nBytes + HEAD);
		}
		else
			nBytes = 0;
	}
	return nBytes;
}

int SocketConnection::sendToKernel(const void * data, unsigned int len )
{
	int ret = 0;
	if ( bitmask & COMPLETE_WRITE )
	{
		bitmask &= ~COMPLETE_WRITE;
		goto pollForWrite;
	}
	ret = TEMP_FAILURE_RETRY( ::send(localSock, data, len, MSG_NOSIGNAL) );
	if ( ERROR == ret && EAGAIN == errno )
	{
pollForWrite:
		ret = waitForWrite();
		if ( 1 == ret )
			ret = TEMP_FAILURE_RETRY( ::send(localSock, data, len, MSG_NOSIGNAL) );
		else
			return ret;
	}
	if ( ret > 0 && static_cast<unsigned int>(ret) < len )
		bitmask |= COMPLETE_WRITE;//the kernel send buf is full, not available.
	return ret;
}

int SocketConnection::send(const void * data, unsigned int len )
{
	int ret = TEMP_FAILURE_RETRY( ::send(localSock, data, len, MSG_NOSIGNAL) );
	if ( ret == ERROR && errno == EAGAIN )
		ret = 0;
	else if ( ret > 0 && static_cast<unsigned int>(ret) < len )
		bitmask |= COMPLETE_WRITE;

	return ret;
}

bool SocketConnection::sendToKernelFully(const void * data, unsigned int len)
{
	unsigned int offset = 0;
	while ( offset < len )
	{
		int nBytes = sendToKernel( (char *)data+offset, len-offset );

		if ( nBytes == ERROR )
			return false;

		offset += nBytes;
	}
	return offset == len;
}

int SocketConnection::sendPacket(const void * data, unsigned int len, bool buffer)
{
	if ( data == NULL || len > MAX_DATASIZE)
		return false;

	bool ret = true;
	StackBuffer buf;
	packPacket(data, buf, len);
	if ( buffer )
	{
		mutex.lock();
		sendBuf.put(buf.readAddr(), buf.sizeForRead());
		mutex.unlock();
	}
	else
	{
		mutex.lock(); //Because linux ::send is not thread safe!
		ret = sendToKernelFully(buf.readAddr(), buf.sizeForRead());
		mutex.unlock();
	}
	return ret;
}

bool SocketConnection::sync()
{
	mutex.lock();
	if ( sendBuf.sizeForRead() > 0 )
	{
		int nBytes = this->send(sendBuf.readAddr(), sendBuf.sizeForRead());
		if ( nBytes >= 0 )
		{
			sendBuf.adjustReadPtr(nBytes);
			mutex.unlock();
			return true;
		}
		else 
		{
			mutex.unlock();
			return false;
		}
	}
	else
	{
		mutex.unlock();
	}
	return true;
}

bool SocketConnection::finalSync()
{
	mutex.lock();
	if ( sendBuf.sizeForRead() > 0 )
	{
		bool ret = this->sendToKernelFully(sendBuf.readAddr(), sendBuf.sizeForRead());
		sendBuf.reset();
		mutex.unlock();
		return ret;
	}
	mutex.unlock();
	return true;
}

